from fastapi import FastAPI
from pydantic import BaseModel
from services.chat.knowledge import load_documents, create_vector_store
from services.chat.qa_chain import QASystem
from utils.logger import setup_logger

logger = setup_logger()

app = FastAPI()

try:
    # 初始化系统
    logger.info("开始加载文档...")
    texts = load_documents()
    logger.info(f"成功加载文档，共 {len(texts)} 条")

    logger.info("开始创建向量存储...")
    vector_store = create_vector_store(texts)
    logger.info("向量存储创建完成")

    # model_path = "D:/demo/gitee/python/models/chatglm3-6b-local"
    # model_path = "THUDM/chatglm3-6b"
    model_path = "D:\demo\gitee\python\models\chatglm2-6b-int4"
    logger.info(f"开始加载模型: {model_path}")
    qa_system = QASystem(model_path, vector_store)
    logger.info("模型加载完成")

except Exception as e:
    logger.error(f"初始化失败: {str(e)}", exc_info=True)
    raise

class Question(BaseModel):
    text: str

@app.post("/ask")
async def ask_question(question: Question):
    try:
        logger.info(f"收到问题: {question.text}")
        response = qa_system.answer_question(question.text)
        logger.info(f"返回答案: {response}")
        return response
    except Exception as e:
        logger.error(f"处理问题时出错: {str(e)}", exc_info=True)
        return {"error": str(e)}

# 启动服务
def chat_run():
    try:
        import uvicorn
        logger.info("开始启动服务器...")
        uvicorn.run(
            app, 
            host="0.0.0.0", 
            port=8000,
            reload=False,
            workers=1,
            log_level="info"
        )
    except Exception as e:
        logger.error(f"服务器启动失败: {str(e)}", exc_info=True)
        raise